#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# 主启动脚本
# QQ: 34538980@qq.com
# Jekkay Hu, 2013.5.5
# ///////////////////////////////////////////////////////////////////
#                            _ooOoo_                               //
#                           o8888888o                              //
#                           88" . "88                              //
#                           (| ^_^ |)                              //
#                           O\  =  /O                              //
#                        ____/`---'\____                           //
#                      .'  \\|     |//  `.                         //
#                     /  \\|||  :  |||//  \                        //
#                    /  _||||| -:- |||||-  \                       //
#                    |   | \\\  -  /// |   |                       //
#                    | \_|  ''\---/''  |   |                       //
#                    \  .-\__  `-`  ___/-. /                       //
#                  ___`. .'  /--.--\  `. . ___                     //
#                ."" '<  `.___\_<|>_/___.'  >'"".                  //
#              | | :  `- \`.;`\ _ /`;.`/ - ` : | |                 //
#              \  \ `-.   \_ __\ /__ _/   .-` /  /                 //
#        ========`-.____`-.___\_____/___.-`____.-'========         //
#                             `=---='                              //
#        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^        //
#             佛祖保佑       永无BUG        运行正常                //
#     -------------------------------------------------------      //
#               QQ: 34538980@qq.com                                //
#               博客: http://www.easysb.cn                         //
#               Jekkay Hu, 2013.5.5                                //
# ///////////////////////////////////////////////////////////////////
#
import logging
import queue
import time
from conf.context import context
from scan.search_keywords_thread import SearchKeywordsThread
from scan.domain_supplier_thread import DomainSupplierThread
from scan.scan_work_thread import ScanWorkThread
from scan.result_domain_sync_thread import DomainSyncThread
from scan.result_target_sync_thread import TargetSyncThread
from scan.links_processor import LinksProcessor

_DEFAULT_SCAN_THREAD = 5
_DEFAULT_QUEUE_SIZE = 100
_QUEUE_SIZE = _DEFAULT_QUEUE_SIZE
try:
    if context.setting['scan'] and context.setting['scan']['queue_size']:
        _QUEUE_SIZE = max(_DEFAULT_QUEUE_SIZE, int(context.setting['scan']['queue_size']))
except Exception as e:
    pass


class MainScan(object):
    def __init__(self):
        self.domain_queue = queue.Queue(_QUEUE_SIZE)
        self.result_target_queue = queue.Queue(_QUEUE_SIZE)
        self.result_domain_queue = queue.Queue(_QUEUE_SIZE)
        self.threads = []

    def scan_targets(self, *links):
        if not links:
            logging.error("no target specified")
            return
        for link in links:
            if str(link).find("://") < 0:
                link = 'http://' + str(link)
            d = LinksProcessor.convert_link_to_domain(link)
            if not d:
                logging.error('invalid target %s' % link)
                continue
            self.domain_queue.put(d.__dict__)
        # self._create_search_thread()
        # self._create_domain_supplier_thread()
        self._create_scan_thread_pool()
        self._create_domain_sync_thread()
        self._create_target_sync_thread()
        self._start_all_thread()
        self._wait_all_thread_finish()

    def start_system(self):
        self._create_search_thread()
        self._create_domain_supplier_thread()
        self._create_scan_thread_pool()
        self._create_domain_sync_thread()
        self._create_target_sync_thread()
        self._start_all_thread()
        self._wait_forever()

    # 启动搜索引擎的线程
    def _create_search_thread(self):
        self.threads.append(SearchKeywordsThread())

    # 启动域名提供线程
    def _create_domain_supplier_thread(self):
        self.threads.append(DomainSupplierThread(self.domain_queue))

    # 工作子线程池
    def _create_scan_thread_pool(self):
        size = _DEFAULT_SCAN_THREAD
        try:
            if context.setting['scan'] and context.setting['scan']['thread_count']:
                size = max(1, int(context.setting['scan']['thread_count']))
        except Exception as e:
            pass
        for i in range(0, size):
            self.threads.append(ScanWorkThread(self.domain_queue, self.result_domain_queue, self.result_target_queue))

    # 域名发现线程
    def _create_domain_sync_thread(self):
        self.threads.append(DomainSyncThread(self.result_domain_queue))

    # 目标发现线程
    def _create_target_sync_thread(self):
        self.threads.append(TargetSyncThread(self.result_target_queue))

    def _start_all_thread(self):
        for thread in self.threads:
            thread.start()

    def _wait_forever(self):
        while True:
            time.sleep(1)

    def _wait_all_thread_finish(self):
        is_any_working = True
        while is_any_working:
            time.sleep(1)
            is_any_working = False
            for thread in self.threads:
                if thread.is_working:
                    is_any_working = True
                    break
        self._terminate_threads()
        logging.info("scan target finish, exit now!")

    def _terminate_threads(self):
        time.sleep(1)
        for thread in self.threads:
            try:
                thread.exit_flag = True
            except Exception as e:
                pass
        time.sleep(1)

